package com.yunyi.socket;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;


@Slf4j
public class MyWebSocketHandler extends TextWebSocketHandler {
    private final StringBuilder response = new StringBuilder();
    private final CountDownLatch latch = new CountDownLatch(1);
    private volatile boolean isCompleted = false;

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String payload = message.getPayload();
        response.append(payload);
        log.info("收到服务端文本消息: {}", payload); // 关键日志

        JSONObject jsonObject = JSON.parseObject(payload);
        if (jsonObject.containsKey("is_final") && !jsonObject.getBoolean("is_final")) {
            log.info("检测到服务端最终响应");
            isCompleted = true;
            latch.countDown();
        }
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        log.warn("收到未处理的二进制消息，大小: {} bytes", message.getPayloadLength());
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("成功连接服务端: {}", session.getRemoteAddress());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        if (!isCompleted) {
            log.error("连接意外关闭，状态码: {}", status.getCode());
            latch.countDown();
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        log.error("传输错误", exception);
        latch.countDown();
    }

    public String getResponse() throws InterruptedException {
        // 等待最大30秒
        if (latch.await(30, TimeUnit.SECONDS)) {
            return response.toString();
        }
        throw new RuntimeException("等待响应超时");
    }
}
